"""
A colllection of wrapper methods around iptables.
Makes constructing tests that handle network disconnection a breeze.
"""
from subprocess import Popen

from storageConnectionTestUtils import eStorageType

_baseCmd = ['sudo', '/sbin/iptables']
NFS_PORT = 2049

def flush():
    _exec(_baseCmd + ['--flush'])

def blockConnection(connection):
    name, type, con = connection
    if type == eStorageType.nfs:
        _exec(_baseCmd + ['-A', 'OUTPUT', '-d', con['connection'].split(":")[0], '-p', 'tcp', '--dport', str(NFS_PORT), '-j', 'DROP'])
    else:
        _exec(_baseCmd + ['-A', 'OUTPUT', '-d', con['connection']]) #TODO : Block specific iscsi port, Handle FCP somehow
def unblockConnection(connection):
    name, type, con = connection
    if type == eStorageType.nfs:
        _exec(_baseCmd + ['-D', 'OUTPUT', '-d', con['connection'].split(":")[0], '-p', 'tcp', '--dport', str(NFS_PORT), '-j', 'DROP'])
    else:
        _exec(_baseCmd + ['-D', 'OUTPUT', '-d', con['connection']]) #TODO : Block specific iscsi port, Handle FCP somehow

def _exec(cmd):
    p = Popen(cmd)
    out, err = p.communicate()
    rc = p.returncode
    if rc != 0:
        raise Exception("iptables command failed (%s)" % cmd)
